import torch
import sys

def check_gpu():
    print(f"Python 版本: {sys.version}")
    print(f"PyTorch 版本: {torch.__version__}")
    print(f"CUDA 是否可用: {torch.cuda.is_available()}")
    
    if torch.cuda.is_available():
        print(f"CUDA 版本: {torch.version.cuda}")
        print(f"当前 GPU: {torch.cuda.get_device_name(0)}")
        print(f"GPU 数量: {torch.cuda.device_count()}")
    else:
        print("警告: CUDA 不可用，请检查:")
        print("1. NVIDIA 驱动是否正确安装")
        print("2. CUDA Toolkit 是否正确安装")
        print("3. PyTorch 是否安装了 CUDA 版本")

if __name__ == "__main__":
    check_gpu() 